#!/usr/bin/env python
# -*- coding: utf-8 -*-

import numpy as np

"""
    @author: yubo.eth

    @created on: 2022/9/3

    @desc: 指标计算工具类
    单独使用请执行 pip3 install numpy
 
    其他指标可参考 python stockstats 库
"""

class IndicatorsUtils(object):


    """
        EMA算法
        cps:价格集合 一般计算ema都以close价格计算
        days:均线天数 eg:days=5则表示是ema5日均线
        return:返回cps相同数量的ema数据
     """
    def calc_ema(cps,days):
        emas = cps.copy()  # 创造一个和cps一样大小的集合
        for i in range(len(cps)):
            if i == 0:
                emas[i] = cps[i]
            if i > 0:
                emas[i] = ((days - 1) * emas[i - 1] + 2 * cps[i]) / (days + 1)
        return emas




    """
        TD序列算法
        kline_data:调用calc_kine_close方法
        return: td数值
    """
    def calc_td_sequence(kline_data):
        close_np = IndicatorsUtils.calc_kine_close(kline_data)
        close_shift = np.empty_like(close_np)
        close_shift[:4] = 0
        close_shift[4:] = close_np[:-4]
        compare_array = close_np > close_shift
        result = np.empty(len(close_np), int)
        counting_number: int = 0
        for i in range(len(close_np)):
            if np.isnan(close_shift[i]):
                result[i] = 0
            else:
                compare_bool = compare_array[i]
                if compare_bool:
                    if counting_number >= 0:
                        counting_number += 1
                    else:
                        counting_number = 1
                else:
                    if counting_number <= 0:
                        counting_number -= 1
                    else:
                        counting_number = -1
                result[i] = counting_number
        return result[-5:]   # 校验数据，取最后5条K线对应的td值


    """
        td算法使用
        kline 算法
        获取kline close价格集合
    """
    def calc_kine_close(klines):
        close_kline_price = []
        for klineclose in klines:
            close_kline_price.append(float(klineclose[4]))
        return close_kline_price